import json
import os
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics.pairwise import cosine_similarity
from typing import Dict, List, Any, Tuple, Optional
import gc

# 设置中文字体支持
# plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC", "Arial Unicode MS"]
plt.rcParams["axes.unicode_minus"] = False  # 解决负号显示问题

class EmbeddingVisualizer:
    """Embeddings向量相关性可视化工具"""
    
    def __init__(self, cache_path: str):
        """初始化可视化工具"""
        self.cache_path = cache_path
        self.pdf_cache = self._load_cache()
        self.similarity_matrix = None
        self.labels = None
        self.document_groups = self._get_document_groups()  # 记录每个embedding属于哪个文档
        
    def _load_cache(self) -> Dict[str, Any]:
        """加载JSON缓存文件"""
        try:
            with open(self.cache_path, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception as e:
            print(f"加载缓存失败: {e}")
            return {}
    
    def _get_document_groups(self) -> List[str]:
        """获取每个embedding所属的文档，用于判断是否为不同文档"""
        groups = []
        for pdf_hash, data in self.pdf_cache.items():
            if 'embeddings' not in data:
                continue
                
            pdf_path = data.get('path', pdf_hash)
            filename = os.path.basename(pdf_path)
            filename = os.path.splitext(filename)[0]
            
            # 为每个embedding记录所属文档
            for _ in data['embeddings']:
                groups.append(filename)
        return groups
    
    def process_embeddings(self) -> Tuple[np.ndarray, List[str]]:
        """处理embeddings数据，生成相似度矩阵和标签"""
        all_embeddings = []
        labels = []
        
        # 提取所有embeddings和对应的标签
        for pdf_hash, data in self.pdf_cache.items():
            if 'embeddings' not in data:
                continue
                
            pdf_path = data.get('path', pdf_hash)
            # 从路径中提取文件名
            filename = os.path.basename(pdf_path)
            # 移除文件扩展名
            filename = os.path.splitext(filename)[0]
            
            # 为每个embedding创建标签
            for i, emb in enumerate(data['embeddings']):
                # 提取前20个字符作为文件名缩写
                short_filename = filename[:20] if len(filename) > 20 else filename
                label = f"{short_filename}_{i}"
                labels.append(label)
                all_embeddings.append(emb)
        
        if not all_embeddings:
            raise ValueError("没有找到有效的embeddings数据")
            
        # 转换为numpy数组
        embeddings_array = np.array(all_embeddings)
        
        # 计算余弦相似度矩阵
        self.similarity_matrix = cosine_similarity(embeddings_array)
        self.labels = labels
        
        # 清理内存
        del embeddings_array
        gc.collect()
        
        return self.similarity_matrix, self.labels
    
    def plot_heatmap(self, output_path: Optional[str] = None, 
                    figsize: Tuple[int, int] = (15, 12),
                    dpi: int = 300,
                    title: Optional[str] = None,
                    similarity_threshold: float = 0.6,
                    different_doc_color: str = 'red',
                    same_doc_color: str = 'blue') -> None:
        """
        绘制相似度热力图，标记不同文档和相同文档的高相似度对
        
        Args:
            output_path: 图像保存路径
            figsize: 图像尺寸
            dpi: 分辨率
            title: 标题
            similarity_threshold: 相似度阈值
            different_doc_color: 不同文档高相似度标记颜色
            same_doc_color: 相同文档高相似度标记颜色
        """
        if self.similarity_matrix is None or self.labels is None:
            self.process_embeddings()
        
        # 创建图形
        plt.figure(figsize=figsize, dpi=dpi)
        
        # 设置标题
        if title is None:
            title = f"Embeddings Cosine Sim (thres: >{similarity_threshold})"
        plt.title(title, fontsize=15)
        
        # 绘制热力图（下三角）
        cmap = sns.diverging_palette(220, 20, as_cmap=True)
        mask = np.triu(np.ones_like(self.similarity_matrix, dtype=bool))  # 遮盖上三角
        
        ax = sns.heatmap(self.similarity_matrix, mask=mask, cmap=cmap, annot=False,
                    square=True, linewidths=.5, cbar_kws={"shrink": .5},
                    xticklabels=self.labels, yticklabels=self.labels)
        
        # 标记高相似度对
        n = len(self.similarity_matrix)
        for i in range(n):
            for j in range(i):  # 下三角区域 i > j
                if self.similarity_matrix[i, j] > similarity_threshold:
                    # 判断是否为相同文档
                    if self.document_groups[i] == self.document_groups[j]:
                        # 相同文档：蓝色方框
                        color = same_doc_color
                    else:
                        # 不同文档：红色方框
                        color = different_doc_color
                    
                    rect = plt.Rectangle((j, i), 1, 1, fill=False, 
                                       edgecolor=color, linewidth=2, zorder=3)
                    ax.add_patch(rect)
        
        # 标签旋转
        plt.xticks(rotation=45, ha='right', fontsize=8)
        plt.yticks(fontsize=8)
        
        plt.tight_layout()
        
        # 保存或显示
        if output_path:
            plt.savefig(output_path, bbox_inches='tight')
            print(f"热力图已保存到 {output_path}")
        else:
            plt.show()
        
        plt.close()
        gc.collect()
    
    def plot_clustermap(self, output_path: Optional[str] = None,
                        figsize: Tuple[int, int] = (15, 12),
                        dpi: int = 300,
                        title: Optional[str] = None,
                        similarity_threshold: float = 0.6,
                        different_doc_color: str = 'red',
                        same_doc_color: str = 'blue') -> None:
        """
        绘制聚类热力图（结构与cos sim一致，不重排），标记不同类型高相似度对
        """
        if self.similarity_matrix is None or self.labels is None:
            self.process_embeddings()
        
        # 创建聚类热力图（关闭行列聚类，保持原始顺序）
        g = sns.clustermap(
            self.similarity_matrix, 
            row_cluster=False,  # 不聚类行
            col_cluster=False,  # 不聚类列
            method='average',
            metric='cosine',
            cmap="YlGnBu", 
            annot=False, 
            figsize=figsize,
            xticklabels=self.labels, 
            yticklabels=self.labels,
            mask=np.triu(np.ones_like(self.similarity_matrix, dtype=bool)),  # 遮盖上三角
            dendrogram_ratio=(0.01, 0.01)  # 隐藏树状图
        )
        
        # 设置标题
        if title is None:
            title = f"Embeddings Cluster Sim (阈值: >{similarity_threshold})"
        g.fig.suptitle(title, fontsize=15, y=0.98)
        
        # 标记高相似度对（因关闭聚类，行列顺序与原始一致）
        n = len(self.similarity_matrix)
        for i in range(n):
            for j in range(i):  # 下三角区域 i > j
                if self.similarity_matrix[i, j] > similarity_threshold:
                    if self.document_groups[i] == self.document_groups[j]:
                        color = same_doc_color
                    else:
                        color = different_doc_color
                    
                    # 在聚类热力图上绘制标记
                    g.ax_heatmap.add_patch(plt.Rectangle((j, i), 1, 1, 
                                                       fill=False, edgecolor=color, 
                                                       linewidth=2, zorder=3))
        
        # 标签旋转
        plt.setp(g.ax_heatmap.get_xticklabels(), rotation=45, ha='right', fontsize=8)
        plt.setp(g.ax_heatmap.get_yticklabels(), fontsize=8)
        
        plt.tight_layout()
        
        # 保存或显示
        if output_path:
            g.savefig(output_path, dpi=dpi, bbox_inches='tight')
            print(f"聚类热力图已保存到 {output_path}")
        else:
            plt.show()
        
        plt.close()
        gc.collect()


def visualize_embeddings(cache_path: str, output_dir: str = '.', 
                        plot_type: str = 'heatmap', **kwargs) -> None:
    """可视化embeddings向量相关性的主函数"""
    os.makedirs(output_dir, exist_ok=True)
    visualizer = EmbeddingVisualizer(cache_path)
    
    if plot_type == 'heatmap':
        output_path = os.path.join(output_dir, 'embeddings_heatmap.png')
        visualizer.plot_heatmap(output_path=output_path,** kwargs)
    elif plot_type == 'clustermap':
        output_path = os.path.join(output_dir, 'embeddings_clustermap.png')
        visualizer.plot_clustermap(output_path=output_path, **kwargs)
    else:
        raise ValueError(f"不支持的图表类型: {plot_type}")


if __name__ == "__main__":
    cache_path = "./cache/embeddings_cache.json"  # 替换为你的缓存路径
    output_dir = "./test_output"
    threshold = 0.6
    
    # 绘制余弦相似度热力图
    visualize_embeddings(
        cache_path=cache_path,
        output_dir=output_dir,
        plot_type='heatmap',
        figsize=(20, 18),
        title=f"PDF Embeddings Cosine Sim (thres: {threshold})",
        similarity_threshold=threshold,
        different_doc_color='red',
        same_doc_color='blue'
    )
    
    # 绘制聚类热力图（结构与cos sim一致）
    visualize_embeddings(
        cache_path=cache_path,
        output_dir=output_dir,
        plot_type='clustermap',
        figsize=(20, 18),
        title=f"PDF Embeddings Cluster Sim (thres: {threshold})",
        similarity_threshold=threshold,
        different_doc_color='red',
        same_doc_color='blue'
    )